"""
权重重排序

相对document进行去重处理
使用BM25算法对query和每个doucment进行分词处理，并计算每个文档和query的相似度分数。这样就得到每个document和query的相似度分数。
对query和document进行向量化处理，并计算每个document和query向量的相似度分数。这样就得到每个document和query的向量的相似度分数。
针对每个document：将基于分词和向量的两个分数按照给定的权重综合计算，得到每个document和query相似度的最终评分。
根据document的得分对其进行排序，并选择得分最高的top_n的个document返回。

"""
import math
from collections import Counter

import numpy as np

from util.T2llm import T2llm
from util.rag.jieba_keyword_handler import JiebaKeywordTableHandler


# 计算关键词相似度
def _calculate_keyword_score(query: str, documents: list[str]) -> list[float]:
        """
        Calculate BM25 scores
        :param query: search query
        :param documents: documents for reranking
        """
        # query 的关键词
        keyword_table_handler = JiebaKeywordTableHandler()
        query_keywords = keyword_table_handler.extract_keywords(query, None)

        # 文档关键词
        documents_keywords = []
        for document in documents:
            # get the document keywords
            document_keywords = keyword_table_handler.extract_keywords(document, None)
            documents_keywords.append(document_keywords)


        # Counter query keywords(TF)
        query_keyword_counts = Counter(query_keywords)

        # total documents
        total_documents = len(documents)

        # calculate all documents' keywords IDF
        all_keywords = set()
        for document_keywords in documents_keywords:
            all_keywords.update(document_keywords)


        # 通过TF-IDF算法计算keyword的评分（重要程度）
        keyword_idf = {}
        for keyword in all_keywords:
            # calculate include query keywords' documents
            # 计算文档中包含该关键字的数量
            doc_count_containing_keyword = sum(1 for doc_keywords in documents_keywords if keyword in doc_keywords)
            # IDF
            keyword_idf[keyword] = math.log((1 + total_documents) / (1 + doc_count_containing_keyword)) + 1


        # 遍历查询字符串计算 每个 关键词的总分
        query_tfidf = {}
        for keyword, count in query_keyword_counts.items():
            tf = count
            idf = keyword_idf.get(keyword, 0)
            query_tfidf[keyword] = tf * idf

        # calculate all documents' TF-IDF
        # 文档总分计算
        documents_tfidf = []
        for document_keywords in documents_keywords:
            document_keyword_counts = Counter(document_keywords)
            document_tfidf = {}
            for keyword, count in document_keyword_counts.items():
                tf = count
                idf = keyword_idf.get(keyword, 0)
                document_tfidf[keyword] = tf * idf
            documents_tfidf.append(document_tfidf)

        # 相似度计算 余弦相似度
        def cosine_similarity(vec1, vec2):
            intersection = set(vec1.keys()) & set(vec2.keys())
            numerator = sum(vec1[x] * vec2[x] for x in intersection)

            sum1 = sum(vec1[x] ** 2 for x in vec1)
            sum2 = sum(vec2[x] ** 2 for x in vec2)
            denominator = math.sqrt(sum1) * math.sqrt(sum2)

            if not denominator:
                return 0.0
            else:
                return float(numerator) / denominator

        similarities = []
        for document_tfidf in documents_tfidf:
            similarity = cosine_similarity(query_tfidf, document_tfidf)
            similarities.append(similarity)

        # for idx, similarity in enumerate(similarities):
        #     print(f"Document {idx + 1} similarity: {similarity}")

        return similarities


def run(
        query: str,
        documents: list[str],
        score_threshold:float= None,
        top_n: int = None,
        keyword_weight: float = 0.7,  #关键词权重 词义
        vector_weight: float = 0.3,   #向量权重 语义
        query_vector_scores: list[float] = None  #文档的向量分数，如果存在向量数据库，查询的时候已经带上了就可以直接传入，不重新计算
) -> list[tuple]: #(document,score)  //文档和分数 排序之后的 分数范围到-1/1之间
    # 对query和每个文档分块进行分词处理，并计算每个文档和query的相似度分数。
    query_scores = _calculate_keyword_score(query, documents)

    # 对query和document进行向量化处理，并计算每个document和query向量的相似度分数
    if query_vector_scores is None:
        llm = T2llm()
        query_vector_scores = llm.calculate_cosine(query, documents)

    rerank_documents = []
    # 针对每个document：将基于分词和向量的两个分数按照给定的权重综合计算，得到最终的评分
    for document, query_score, query_vector_score in zip(documents, query_scores, query_vector_scores):
        # format document
        # 计算综合得分
        score = (
                vector_weight * query_vector_score
                + keyword_weight * query_score
        )
        # 如果设置了 score_threshold 且文档的综合评分为负数，则跳过
        if score_threshold and score < score_threshold:
            continue
        # 添加文档的综合评分到元数据中
        rerank_documents.append((document, score))

        # 将所有符合要求的文档按分数从高到低排序。
    rerank_documents.sort(key=lambda x: x[1], reverse=True)

    # 根据 top_n 参数返回评分最高的前 N 个文档；如果不设置 top_n，则返回所有符合条件的文档。
    return rerank_documents[:top_n] if top_n else rerank_documents

if __name__ == '__main__':
    query = "苹果"
    documents = ["苹果","大苹果","香蕉","梨子"]
    result = run(query,documents,top_n=3)
    print(result)
    # [('苹果', np.float64(0.9999999999999999)), ('大苹果', np.float64(0.9999999999999999)), ('香蕉', np.float64(0.29999999999999993))]
